#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import time
import subprocess

java_sdk_path = "/root/KWDB-sdk/java-sdk"
server_path = "/root/KWDB-sdk/server"

def check_server():
    """检查服务器是否运行"""
    try:
        result = subprocess.run(["netstat", "-tlnp"], capture_output=True, text=True, timeout=5)
        if "9090" in result.stdout:
            print("✓ 服务器正在运行（端口9090）")
            return True
        result = subprocess.run(["ss", "-tlnp"], capture_output=True, text=True, timeout=5)
        if "9090" in result.stdout:
            print("✓ 服务器正在运行（端口9090）")
            return True
        print("✗ 服务器未运行")
        return False
    except:
        return False

def start_server():
    """启动服务器"""
    print("\n=== 启动服务器 ===")
    try:
        os.chdir(server_path)
        server_process = subprocess.Popen(
            ["java", "-jar", "target/kwdb-server-1.0.0.jar", "--port", "9090"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        time.sleep(8)
        if server_process.poll() is None:
            print("✓ 服务器启动成功")
            return server_process
        else:
            print("✗ 服务器启动失败")
            return None
    except Exception as e:
        print(f"✗ 服务器启动异常: {e}")
        return None

def run_performance_test():
    """运行性能测试"""
    print("\n=== 性能测试 ===")
    
    perf_test_code = """
import com.kwdb.sdk.KWDBClient;
import com.kwdb.sdk.config.ClientConfig;
import com.kwdb.sdk.data.DataPoint;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class PerformanceTest {
    public static void main(String[] args) {
        try {
            ClientConfig config = new ClientConfig("localhost", 9090)
                    .setBatchSize(1000)
                    .setBatchTimeout(100)
                    .setEnableCompression(true);
            
            KWDBClient client = new KWDBClient(config);
            client.connect();
            
            System.out.println("=== 性能测试开始 ===");
            
            // 测试1: 单线程写入性能
            int recordCount = 10000;
            List<DataPoint> dataPoints = new ArrayList<>();
            
            for (int i = 0; i < recordCount; i++) {
                DataPoint dp = new DataPoint("perf_test")
                        .addTag("thread", "single")
                        .addTag("id", String.valueOf(i))
                        .addField("value", Math.random() * 100)
                        .addField("timestamp", System.currentTimeMillis());
                dataPoints.add(dp);
            }
            
            long startTime = System.currentTimeMillis();
            client.writeBatch(dataPoints);
            long endTime = System.currentTimeMillis();
            
            long duration = endTime - startTime;
            double throughput = (double) recordCount / duration * 1000;
            
            System.out.println("\\n单线程写入性能:");
            System.out.println("记录数: " + recordCount);
            System.out.println("耗时: " + duration + "ms");
            System.out.println("吞吐量: " + String.format("%.2f", throughput) + " 记录/秒");
            
            // 测试2: 多线程写入性能
            int threadCount = 10;
            int recordsPerThread = 1000;
            ExecutorService executor = Executors.newFixedThreadPool(threadCount);
            
            List<CompletableFuture<Void>> futures = new ArrayList<>();
            
            startTime = System.currentTimeMillis();
            
            for (int t = 0; t < threadCount; t++) {
                final int threadId = t;
                CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                    try {
                        List<DataPoint> threadData = new ArrayList<>();
                        for (int i = 0; i < recordsPerThread; i++) {
                            DataPoint dp = new DataPoint("perf_test")
                                    .addTag("thread", "multi")
                                    .addTag("thread_id", String.valueOf(threadId))
                                    .addTag("id", String.valueOf(i))
                                    .addField("value", Math.random() * 100)
                                    .addField("timestamp", System.currentTimeMillis());
                            threadData.add(dp);
                        }
                        client.writeBatch(threadData);
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }, executor);
                futures.add(future);
            }
            
            CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
            
            endTime = System.currentTimeMillis();
            duration = endTime - startTime;
            int totalRecords = threadCount * recordsPerThread;
            throughput = (double) totalRecords / duration * 1000;
            
            System.out.println("\\n多线程写入性能:");
            System.out.println("线程数: " + threadCount);
            System.out.println("总记录数: " + totalRecords);
            System.out.println("耗时: " + duration + "ms");
            System.out.println("吞吐量: " + String.format("%.2f", throughput) + " 记录/秒");
            
            executor.shutdown();
            client.disconnect();
            
            System.out.println("\\n=== 性能测试完成 ===");
            
        } catch (Exception e) {
            System.err.println("性能测试失败: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
"""
    
    perf_file = os.path.join(java_sdk_path, "PerformanceTest.java")
    with open(perf_file, "w", encoding="utf-8") as f:
        f.write(perf_test_code)
    
    try:
        os.chdir(java_sdk_path)
        
        # 编译
        result = subprocess.run(
            ["javac", "-cp", "target/classes:target/dependency/*", "PerformanceTest.java"],
            capture_output=True, text=True, timeout=60
        )
        
        if result.returncode != 0:
            print(f"✗ 编译失败: {result.stderr}")
            return False
        
        # 运行
        result = subprocess.run(
            ["java", "-cp", ".:target/classes:target/dependency/*", "PerformanceTest"],
            capture_output=True, text=True, timeout=120
        )
        
        if result.returncode == 0:
            print("✓ 性能测试完成")
            print(result.stdout)
            return True
        else:
            print(f"✗ 性能测试失败: {result.stderr}")
            return False
            
    except Exception as e:
        print(f"✗ 性能测试异常: {e}")
        return False

if __name__ == "__main__":
    print("=== KWDB 性能测试 ===")
    
    server_process = None
    if not check_server():
        server_process = start_server()
        if not server_process:
            sys.exit(1)
    
    try:
        success = run_performance_test()
        sys.exit(0 if success else 1)
    finally:
        if server_process:
            server_process.terminate()
            server_process.wait()
            print("\n服务器已停止")

